import numpy as np
import librosa
import matplotlib.pyplot as plt
from scipy.signal import butter, filtfilt
import noisereduce as nr
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 或 ['Microsoft YaHei']，需系统已安装
matplotlib.rcParams['axes.unicode_minus'] = False    # 正确显示负号

def calculate_snr(audio_path):
    """
    计算MP3音频文件的语音信噪比(SNR)
    """
    # 加载音频文件
    y, sr = librosa.load(audio_path, sr=16000, mono=True)
    
    # 预处理：带通滤波(300-3000Hz)
    nyq = 0.5 * sr
    low = 300 / nyq
    high = 3000 / nyq
    b, a = butter(4, [low, high], btype='band')
    y = filtfilt(b, a, y)
    
    # 使用noisereduce进行噪声估计
    # 使用前100ms作为噪声样本
    noise_sample = y[:int(0.1 * sr)] if len(y) > int(0.1 * sr) else y
    
    # 应用噪声减少（但不真正减少，只用于噪声估计）
    reduced_noise = nr.reduce_noise(
        y=y, 
        y_noise=noise_sample,
        sr=sr,
        stationary=True
        # 不设置 prop_decrease=0
    )
    
    # 计算噪声分量
    noise = y - reduced_noise
    
    # 计算信号功率
    signal_power = np.sum(reduced_noise ** 2) / len(reduced_noise)
    
    # 计算噪声功率
    noise_power = np.sum(noise ** 2) / len(noise)
    
    # 避免除零错误
    if noise_power < 1e-10:
        return float('inf')
    
    # 计算SNR (dB)
    snr_db = 10 * np.log10(signal_power / noise_power)
    
    # 可视化结果
    plot_snr_analysis(y, reduced_noise, noise, sr, snr_db)
    
    return snr_db

def plot_snr_analysis(original, signal, noise, sr, snr_db):
    """可视化SNR分析结果"""
    plt.figure(figsize=(15, 10))
    
    # 1. 原始音频波形
    plt.subplot(3, 1, 1)
    time = np.arange(len(original)) / sr
    plt.plot(time, original, label='原始音频', alpha=0.7)
    plt.title(f'音频波形 (SNR: {snr_db:.2f} dB)')
    plt.xlabel('时间 (秒)')
    plt.ylabel('幅度')
    plt.legend()
    
    # 2. 分离的信号和噪声
    plt.subplot(3, 1, 2)
    plt.plot(time, signal, 'g', label='语音信号', alpha=0.7)
    plt.plot(time, noise, 'r', label='背景噪声', alpha=0.5)
    plt.title('语音信号与背景噪声分离')
    plt.xlabel('时间 (秒)')
    plt.ylabel('幅度')
    plt.legend()
    
    # 3. 频谱分析
    plt.subplot(3, 1, 3)
    S_original = librosa.amplitude_to_db(np.abs(librosa.stft(original)), ref=np.max)
    S_signal = librosa.amplitude_to_db(np.abs(librosa.stft(signal)), ref=np.max)
    S_noise = librosa.amplitude_to_db(np.abs(librosa.stft(noise)), ref=np.max)
    
    plt.plot(librosa.fft_frequencies(sr=sr, n_fft=2048), 
             np.mean(S_original, axis=1), 'b', label='原始音频频谱', alpha=0.7)
    plt.plot(librosa.fft_frequencies(sr=sr, n_fft=2048), 
             np.mean(S_signal, axis=1), 'g', label='语音信号频谱', alpha=0.7)
    plt.plot(librosa.fft_frequencies(sr=sr, n_fft=2048), 
             np.mean(S_noise, axis=1), 'r', label='背景噪声频谱', alpha=0.7)
    
    plt.title('平均频谱分析')
    plt.xlabel('频率 (Hz)')
    plt.ylabel('幅度 (dB)')
    plt.xlim(0, 4000)  # 限制在语音相关频率范围
    plt.legend()
    
    plt.tight_layout()
    plt.savefig('snr_analysis.png')
    plt.close()
    print("SNR分析图已保存为 'snr_analysis.png'")

def interpret_snr(snr_db):
    """解释SNR值的含义"""
    if snr_db > 40:
        return "极高质量: 背景噪声几乎不可闻"
    elif snr_db > 30:
        return "高质量: 背景噪声轻微，不影响语音理解"
    elif snr_db > 20:
        return "中等质量: 可听到背景噪声但不影响主要语音"
    elif snr_db > 10:
        return "较低质量: 背景噪声明显，可能影响语音理解"
    elif snr_db > 0:
        return "低质量: 背景噪声强烈，语音理解困难"
    else:
        return "极低质量: 背景噪声主导，语音几乎不可理解"

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("用法: python snr_calculator.py <音频文件路径>")
        sys.exit(1)
    
    audio_file = sys.argv[1]
    
    print(f"分析音频文件: {audio_file}")
    snr = calculate_snr(audio_file)
    
    print("\n" + "="*50)
    print(f"语音信噪比(SNR)分析结果:")
    print(f"SNR值: {snr:.2f} dB")
    print(f"质量评估: {interpret_snr(snr)}")
    print("="*50)
    print("提示:")
    print("- SNR > 30 dB 表示高质量的语音录音")
    print("- SNR < 20 dB 表示有明显的背景噪声")
    print("- SNR < 10 dB 表示语音可能难以理解")
    print("详细分析图已保存为 'snr_analysis.png'")